#!/usr/bin/env python3
"""
简单的HLS密钥服务器
用于提供加密密钥给HLS播放器
"""

import os
import json
import hmac
import hashlib
import argparse
from pathlib import Path
from flask import Flask, Response, request, jsonify
from flask_cors import CORS

def load_config(config_file='config.json'):
    """加载配置文件"""
    try:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        config_path = os.path.join(script_dir, config_file)

        if not os.path.exists(config_path):
            raise ValueError(f"配置文件不存在: {config_path}")

        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)

        # 验证必要的配置项
        required_configs = {
            'common.output_dir': config.get('common', {}).get('output_dir'),
            'key_server.host': config.get('key_server', {}).get('host'),
            'key_server.port': config.get('key_server', {}).get('port'),
        }

        missing_configs = [key for key, value in required_configs.items() if value is None]
        if missing_configs:
            raise ValueError(f"配置文件缺少必要参数: {', '.join(missing_configs)}")

        return config
    except Exception as e:
        print(f"❌ 读取配置文件失败: {e}")
        raise

# 加载全局配置
try:
    config = load_config()
    OUTPUT_DIR = config['common']['output_dir']
    KEY_SERVER_HOST = config['key_server']['host']
    KEY_SERVER_PORT = config['key_server']['port']
    ENABLE_CORS = config['key_server'].get('cors', True)
    ENABLE_DEBUG = config['common'].get('enable_debug', False)
except Exception as e:
    print("❌ 初始化失败：请检查配置文件")
    raise SystemExit(1)

print("✅ 配置加载成功:")
print(f"   输出目录: {OUTPUT_DIR}")
print(f"   服务器地址: {KEY_SERVER_HOST}:{KEY_SERVER_PORT}")
print(f"   CORS: {'启用' if ENABLE_CORS else '禁用'}")
print(f"   调试模式: {'启用' if ENABLE_DEBUG else '禁用'}")

app = Flask(__name__)
if ENABLE_CORS:
    CORS(app)

# 全局变量存储密钥信息
key_store = {}

@app.route('/key/<video_id>')
def get_key(video_id):
    """获取视频解密密钥"""
    if video_id not in key_store:
        return Response("Key not found", status=404)

    video_info = key_store[video_id]
    key_hex = video_info['encryption']['key']
    
    # 获取请求头中的签名
    request_signature = request.headers.get('X-Signature', '')
    if not request_signature:
        print(f"❌ 缺少签名: {video_id}")
        return Response("Signature required", status=401)

    # 验证请求签名
    config = load_config()
    secret_key = config.get('signature', {}).get('secret_key', '')
    if not secret_key:
        print("❌ 未配置签名密钥")
        return Response("Signature key not configured", status=500)
    
    # 获取请求URL和密钥服务器信息
    request_url = request.headers.get('X-Original-URL', '')
    key_server = request.headers.get('X-Key-Server', '')
    
    # 组合签名数据（与App播放器对齐）
    # 注意：App播放器使用 URL|KeyServer|timestamp 格式
    # 我们需要尝试解析时间戳并验证签名
    import time
    
    # 尝试多种可能的签名数据组合
    signature_variants = [
        f"{request_url}|{key_server}|{video_id}",  # 原始方案
        f"{request_url}|{key_server}",             # 简化方案
        f"{key_server}|{video_id}",               # 另一种简化方案
    ]
    
    # 检查签名是否包含时间戳（App播放器格式）
    # 如果请求签名是有效的，我们尝试解析时间戳部分
    valid_signature = False
    expected_signature = ""
    
    for variant in signature_variants:
        expected = hmac.new(secret_key.encode(), variant.encode(), hashlib.sha256).hexdigest()
        if hmac.compare_digest(request_signature, expected):
            valid_signature = True
            expected_signature = expected
            break
    
    if not valid_signature:
        print(f"❌ 签名验证失败: {video_id}")
        print(f"   请求签名: {request_signature}")
        print(f"   期望签名: {expected_signature}")
        print(f"   签名数据尝试了多种组合")
        return Response("Signature verification failed", status=403)

    try:
        key_bytes = bytes.fromhex(key_hex)
    except ValueError:
        return Response("Invalid key format", status=500)

    # 记录访问日志
    client_ip = request.remote_addr
    user_agent = request.headers.get('User-Agent', 'Unknown')
    
    # 打印完整的请求参数
    print(f"🔑 密钥请求: {video_id}")
    print(f"   - 客户端IP: {client_ip}")
    print(f"   - User-Agent: {user_agent}")
    print(f"   - 请求URL: {request_url}")
    print(f"   - 密钥服务器: {key_server}")
    print(f"   - 请求签名: {request_signature}")
    print(f"   - 期望签名: {expected_signature}")
    print(f"   - 签名数据组合: {signature_variants}")
    print(f"✅ 签名验证成功: {video_id}")

    return Response(key_bytes, mimetype='application/octet-stream')

@app.route('/keyinfo/<video_id>')
def get_key_info(video_id):
    """获取密钥信息（调试用）"""
    if video_id not in key_store:
        return jsonify({"error": "Key not found"}), 404
    
    # 创建安全的响应（不包含实际密钥）
    info = key_store[video_id].copy()
    if 'encryption' in info:
        info['encryption'] = {
            'method': info['encryption']['method'],
            'type': info['encryption']['type'],
            'segment_count': info['encryption']['segment_count']
        }
    return jsonify(info)

@app.route('/keys')
def list_keys():
    """列出所有可用的密钥ID"""
    return jsonify(list(key_store.keys()))

@app.route('/health')
def health_check():
    """健康检查"""
    return jsonify({
        "status": "ok",
        "keys_count": len(key_store)
    })

def load_key_from_metadata(metadata_file, video_id):
    """从元数据文件加载密钥信息"""
    try:
        with open(metadata_file, 'r', encoding='utf-8') as f:
            metadata = json.load(f)

        encryption = metadata.get('encryption', {})
        if not encryption:
            print(f"⚠️  警告: {metadata_file} 中没有加密信息")
            return False

        if encryption['type'] != 'single-key':
            print(f"⚠️  警告: 不支持的加密类型: {encryption['type']}")
            return False

        key_store[video_id] = metadata
        print(f"✅ 加载密钥: {video_id} from {metadata_file}")
        print(f"   - 加密方法: {encryption['method']}")
        print(f"   - 密钥URI: {encryption['key_uri']}")
        return True

    except Exception as e:
        print(f"❌ 加载密钥失败 {metadata_file}: {e}")
        return False

def load_keys_from_directory(directory):
    """从目录中加载所有密钥"""
    directory = Path(directory)
    loaded_count = 0
    
    for metadata_file in directory.rglob('metadata.json'):
        # 使用父目录名作为video_id
        video_id = metadata_file.parent.name
        if load_key_from_metadata(metadata_file, video_id):
            loaded_count += 1
    
    print(f"从 {directory} 加载了 {loaded_count} 个密钥")
    return loaded_count

def main():
    parser = argparse.ArgumentParser(description='HLS密钥服务器')
    parser.add_argument('--host', default='0.0.0.0', help='服务器地址')
    parser.add_argument('--port', type=int, default=8080, help='服务器端口')
    parser.add_argument('--keys-dir', help='密钥目录路径')
    parser.add_argument('--metadata-file', help='单个元数据文件路径')
    parser.add_argument('--video-id', help='视频ID（与--metadata-file一起使用）')
    parser.add_argument('--debug', action='store_true', help='调试模式')
    
    args = parser.parse_args()
    
    # 加载密钥
    if args.keys_dir:
        load_keys_from_directory(args.keys_dir)
    elif args.metadata_file and args.video_id:
        if not load_key_from_metadata(args.metadata_file, args.video_id):
            print("加载密钥失败")
            return
    else:
        print("请指定 --keys-dir 或 --metadata-file 和 --video-id")
        return
    
    if not key_store:
        print("没有加载任何密钥")
        return
    
    print(f"密钥服务器启动: http://{args.host}:{args.port}")
    print(f"可用密钥: {list(key_store.keys())}")
    print(f"密钥URL格式: http://{args.host}:{args.port}/key/<video_id>")
    
    app.run(host=args.host, port=args.port, debug=args.debug)

if __name__ == "__main__":
    main()